-- Copyright (c) 2025 Huawei Technologies Co., Ltd.
-- openUBMC is licensed under Mulan PSL v2.
-- You can use this software according to the terms and conditions of the Mulan PSL v2.
-- You may obtain a copy of Mulan PSL v2 at: http://license.coscl.org.cn/MulanPSL2
-- THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
-- EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
-- MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
-- See the Mulan PSL v2 for more details.

local lu        = require 'luaunit'
local log       = require 'mc.logging'
local enum      = require 'class.types.types'
local spdm_main = require 'spdm_core'
local spdm_core = spdm_main.core
local spdm_cb   = spdm_main.callbacks
local common    = require 'common'

local DEFAULT_NAME     = "test"
local DEFAULT_TYPE     = "Discrete"
local DEFAULT_PATH     = "/bmc/kepler/Systems/1/Storage/Controllers/1"
local DEFAULT_REQ_DATA = "123"
local DEFAULT_RSP_DATA = "123456"

local function send_recv_cb(device_name, data, wait_rsp, timeout)
    log:notice("================================ cb is called")
    -- 通过make unit_test可以在控制台打印print日志
    print("\n================================ cb is called")
    lu.assertEquals(device_name, DEFAULT_NAME)
    lu.assertEquals(data, DEFAULT_REQ_DATA)
    lu.assertEquals(wait_rsp, 0)
    lu.assertEquals(timeout, 0)
    local rsp = DEFAULT_RSP_DATA
    return 0, rsp
end

function TestSpdm:test_init_send_recv_cb_by_lua()
    local mctp_callbacks = require 'infrastructure.callback.mctp_callbacks'
    local mctp_lib = require 'mctp_lib'
    -- spdm 初始化
    local spdm_c = spdm_core.new()
    spdm_c:init()

    -- 创建responder
    spdm_c:new_responder(DEFAULT_NAME, DEFAULT_TYPE, DEFAULT_PATH)

    -- 查看responder当前计数，预期1个
    count = spdm_c:get_responder_count()
    lu.assertEquals(count, 1)

    -- 创建上下文
    spdm_c:new_spdm_ctx(DEFAULT_NAME)

    -- 初始化上下文
    spdm_c:init_spdm_ctx(DEFAULT_NAME, enum.TRANSFER_TYPE.MCTP:value())

    -- 打桩替换
    local get_phy_addr = mctp_lib.bdf_to_phy_addr
    local get_ep = mctp_lib.get_endpoint_and_transport

    mctp_lib.bdf_to_phy_addr = function(...)
        return 1
    end
    local mock_ep = {}
    mock_ep.Request = function(...)
        return DEFAULT_RSP_DATA
    end

    mctp_lib.get_endpoint_and_transport = function()
        return mock_ep
    end

    mctp_callbacks.register_dbus({})
    mctp_callbacks.register_mctp_send_recv_cb({device_name = DEFAULT_NAME, dev = {}})

    -- 执行发送行为
    lu.assertEquals(spdm_c:try_send_message_to_cb(DEFAULT_NAME, DEFAULT_REQ_DATA), 0);

    -- 确认收消息
    lu.assertEquals(spdm_c:try_get_message(DEFAULT_NAME), DEFAULT_RSP_DATA)

    -- 还原
    mctp_lib.bdf_to_phy_addr = get_phy_addr
    mctp_lib.get_endpoint_and_transport = get_ep

    -- spdm释放
    spdm_c:free()
end

function TestSpdm:test_init_send_recv_cb()
    -- spdm 初始化
    local spdm_c = spdm_core.new()
    spdm_c:init()

    -- 创建responder
    spdm_c:new_responder(DEFAULT_NAME, DEFAULT_TYPE, DEFAULT_PATH)

    -- 查看responder当前计数，预期1个
    count = spdm_c:get_responder_count()
    lu.assertEquals(count, 1)

    -- 创建上下文
    spdm_c:new_spdm_ctx(DEFAULT_NAME)

    -- 初始化上下文
    spdm_c:init_spdm_ctx(DEFAULT_NAME, enum.TRANSFER_TYPE.MCTP:value())

    -- 注册一个回调函数
    local cbs = spdm_cb.new()
    cbs:set_send_recv_cb(DEFAULT_NAME, send_recv_cb);

    -- 执行发送行为
    lu.assertEquals(spdm_c:try_send_message_to_cb(DEFAULT_NAME, DEFAULT_REQ_DATA), 0);

    -- 确认收消息
    lu.assertEquals(spdm_c:try_get_message(DEFAULT_NAME), DEFAULT_RSP_DATA)

    -- spdm释放
    spdm_c:free()
end

function TestSpdm:test_spdm_init()
    -- spdm 初始化
    local spdm_c = spdm_core.new()
    spdm_c:init()

    -- 查看responder初始计数，预期0个
    local count = spdm_c:get_responder_count()
    lu.assertEquals(count, 0)

    -- 创建responder
    spdm_c:new_responder(DEFAULT_NAME, DEFAULT_TYPE, DEFAULT_PATH)

    -- 查看responder当前计数，预期1个
    count = spdm_c:get_responder_count()
    lu.assertEquals(count, 1)

    -- 创建上下文
    spdm_c:new_spdm_ctx(DEFAULT_NAME)

    -- 初始化上下文
    spdm_c:init_spdm_ctx(DEFAULT_NAME, enum.TRANSFER_TYPE.MCTP:value())

    -- 检查上下文
    lu.assertEquals(spdm_c:check_spdm_ctx(DEFAULT_NAME), 0)

    -- 释放上下文
    spdm_c:free_spdm_ctx(DEFAULT_NAME)

    -- spdm释放
    spdm_c:free()
end

function TestSpdm:test_get_digest()
    -- spdm 初始化
    local spdm_c = spdm_core.new()
    spdm_c:init()

    -- 创建responder
    spdm_c:new_responder(DEFAULT_NAME, DEFAULT_TYPE, DEFAULT_PATH)

    -- 创建上下文
    spdm_c:new_spdm_ctx(DEFAULT_NAME)

    -- 调用GET_DIGEST
    local res, mask = spdm_c:get_digests(DEFAULT_NAME)
    lu.assertEquals(res,  0x0)
    lu.assertEquals(mask, 0x7) -- 打桩预期此处mask应该为 0111 = 7

    -- spdm释放
    spdm_c:free()
end

function TestSpdm:test_get_certificate()
    -- spdm 初始化
    local spdm_c = spdm_core.new()
    spdm_c:init()

    -- 创建responder
    spdm_c:new_responder(DEFAULT_NAME, DEFAULT_TYPE, DEFAULT_PATH)

    -- 创建上下文
    spdm_c:new_spdm_ctx(DEFAULT_NAME)

    -- 初始化上下文
    spdm_c:init_spdm_ctx(DEFAULT_NAME, enum.TRANSFER_TYPE.MCTP:value())

    -- 构造一个打桩的connection
    lu.assertEquals(spdm_c:init_connection(DEFAULT_NAME), 0)

    local cert_info = spdm_c:get_certificate(DEFAULT_NAME, 0)
    lu.assertEquals(cert_info.certificateType, 0) -- 预期证书类型PEM = 0
    lu.assertEquals(cert_info.certificateUsageType, 5) -- 预期证书用途 TrustedComponentCertificate = 5
    lu.assertEquals(cert_info.fingerPrint, "ca035e3805da28424d9fd3c938c7240abc6c52d6a78def1cac5f09c8134a0920")
    lu.assertEquals(cert_info.fingerPrintHashAlgorithm, "TCG_ALG_SHA256")
    lu.assertEquals(cert_info.issuer, "CN=TestCA, OU=IT, O=HW, L=CD, S=SC, C=CN")
    lu.assertEquals(cert_info.subject, "CN=TestCA, OU=IT, O=HW, L=CD, S=SC, C=CN")
    lu.assertEquals(cert_info.keyUsage, 0)
    lu.assertEquals(cert_info.keyUsageString, "")
    lu.assertEquals(cert_info.serialNumber, "6c 7f 0d ff 1f 94 fc c2 12 43 94 00 b5 81 c5 9a 9e 77 74 e1")
    lu.assertEquals(cert_info.signatureAlgorithm, "sha256WithRSAEncryption")
    lu.assertEquals(cert_info.validNotBefore, 1763606155) -- 1763606155 对应时间：2025-11-20 10:35:55
    lu.assertEquals(cert_info.validNotAfter, 3340406155) -- 3340406155 对应时间：2075-11-08 10:35:55
    lu.assertEquals(cert_info.certCount, 1) -- 预期只有1本证书
    lu.assertEquals(cert_info.keyLength, 4096) -- 密钥长度 4096
    lu.assertEquals(cert_info.commonName, "TestCA")
    lu.assertEquals(cert_info.hashVal, "400d7fbc")

    -- spdm释放
    spdm_c:free()
end

function TestSpdm:test_challenge()
    -- spdm 初始化
    local spdm_c = spdm_core.new()
    spdm_c:init()

    -- 创建responder
    spdm_c:new_responder(DEFAULT_NAME, DEFAULT_TYPE, DEFAULT_PATH)

    -- 创建上下文
    spdm_c:new_spdm_ctx(DEFAULT_NAME)

    -- 初始化上下文
    spdm_c:init_spdm_ctx(DEFAULT_NAME, enum.TRANSFER_TYPE.MCTP:value())

    -- 构造一个打桩的connection
    lu.assertEquals(spdm_c:init_connection(DEFAULT_NAME), 0)

    -- 获取一个slot_id为0的证书
    spdm_c:get_certificate(DEFAULT_NAME, 0)

    -- 进行slot_id为0的挑战，成功
    lu.assertEquals(spdm_c:challenge_auth(DEFAULT_NAME, 0), 0);

    -- 进行slot_id为1的挑战，失败
    lu.assertEquals(spdm_c:challenge_auth(DEFAULT_NAME, 1), -1);

    -- 进行slot_id为9的挑战，失败
    lu.assertEquals(spdm_c:challenge_auth(DEFAULT_NAME, 9), -1);

    -- spdm释放
    spdm_c:free()
end

function TestSpdm:test_init_measurement()
    -- spdm 初始化
    local spdm_c = spdm_core.new()
    spdm_c:init()

    -- 创建responder
    spdm_c:new_responder(DEFAULT_NAME, DEFAULT_TYPE, DEFAULT_PATH)

    -- 创建上下文
    spdm_c:new_spdm_ctx(DEFAULT_NAME)

    -- 初始化上下文
    spdm_c:init_spdm_ctx(DEFAULT_NAME, enum.TRANSFER_TYPE.MCTP:value())

    -- 检查上下文
    lu.assertEquals(spdm_c:check_spdm_ctx(DEFAULT_NAME), 0)

    -- 构造一个打桩的connection
    lu.assertEquals(spdm_c:init_connection(DEFAULT_NAME), 0)

    -- 获取一个slot_id为0的证书
    spdm_c:get_certificate(DEFAULT_NAME, 0)

    -- 初始化测量信息成功
    lu.assertEquals(spdm_c:init_measurements(DEFAULT_NAME, 0), 0)

    -- 获取测量集合
    local measurement_set = spdm_c:get_measurement_set(DEFAULT_NAME);
    lu.assertEquals(measurement_set.specification, "DMTF")
    lu.assertEquals(measurement_set.summaryType, "TCB")
    lu.assertEquals(measurement_set.summaryHashAlgorithm, "TPM_ALG_SHA_256")
    lu.assertEquals(measurement_set.summaryHash, "0101010101010101010101010101010101010101010101010101010101010101")
    lu.assertEquals(measurement_set.measurementsCount, 3) -- 预期有3个度量块（打桩）

    -- 获取单条测量信息
    -- 获取第1个度量块（打桩）
    local measurement = spdm_c:get_one_measurement(DEFAULT_NAME, 1);
    lu.assertEquals(measurement.hashAlgorithm, "TPM_ALG_SHA_256")
    lu.assertEquals(measurement.measurementType, "FirmwareConfiguration")
    lu.assertEquals(measurement.measurement, "0000000000000000000000000000000000000000000000000000000000000000")
    lu.assertEquals(measurement.partOfSummaryHash, true)
    lu.assertEquals(measurement.securityVersionNumber, "")

    -- 获取第2个度量块（打桩）
    measurement = spdm_c:get_one_measurement(DEFAULT_NAME, 2);
    lu.assertEquals(measurement.hashAlgorithm, "TPM_ALG_SHA_256")
    lu.assertEquals(measurement.measurementType, "FirmwareConfiguration")
    lu.assertEquals(measurement.measurement, "0000000000000000000000000000000000000000000000000000000000000001")
    lu.assertEquals(measurement.partOfSummaryHash, true)
    lu.assertEquals(measurement.securityVersionNumber, "")

    -- 获取第3个度量块（打桩）
    measurement = spdm_c:get_one_measurement(DEFAULT_NAME, 3);
    lu.assertEquals(measurement.hashAlgorithm, "TPM_ALG_SHA_256")
    lu.assertEquals(measurement.measurementType, "FirmwareConfiguration")
    lu.assertEquals(measurement.measurement, "0000000000000000000000000000000000000000000000000000000000000002")
    lu.assertEquals(measurement.partOfSummaryHash, true)
    lu.assertEquals(measurement.securityVersionNumber, "")

    lu.assertErrorMsgContains("get measurement failed", function()
        -- 获取第4个度量块（不存在）
        spdm_c:get_one_measurement(DEFAULT_NAME, 4)
    end)

    -- spdm释放
    spdm_c:free()
end

function TestSpdm:test_get_signed_measurement()
    -- spdm 初始化
    local spdm_c = spdm_core.new()
    spdm_c:init()

    -- 创建responder
    spdm_c:new_responder(DEFAULT_NAME, DEFAULT_TYPE, DEFAULT_PATH)

    -- 创建上下文
    spdm_c:new_spdm_ctx(DEFAULT_NAME)

    -- 初始化上下文
    spdm_c:init_spdm_ctx(DEFAULT_NAME, enum.TRANSFER_TYPE.MCTP:value())

    -- 构造一个打桩的connection
    lu.assertEquals(spdm_c:init_connection(DEFAULT_NAME), 0)

    -- 获取一个slot_id为0的证书
    spdm_c:get_certificate(DEFAULT_NAME, 0)

    -- 初始化测量信息成功
    lu.assertEquals(spdm_c:init_measurements(DEFAULT_NAME, 0), 0)

    local indices = {1, 2, 3} -- 使用1、2、3三个槽位进行签名
    local result = spdm_c:get_signed_measurement(DEFAULT_NAME, indices, 0,
        "0102030405060708090A0B0C0D0E0F100102030405060708090A0B0C0D0E0F10")

    lu.assertEquals(result.version, "1.0")
    lu.assertEquals(result.hashAlgorithm, "TPM_ALG_SHA_256")
    lu.assertEquals(result.signAlgorithm, "TPM_ALG_RSAPSS_2048")
    lu.assertEquals(result.publicKey, "")

    local expect_measurement = "00000000000000000000000000000000000000000000000000000000000000000000000000000000000" ..
        "0000000000000000000000000000000000000000000010101010101010101010101010101010101010101010101010101010101010101"
    lu.assertEquals(result.measurement, common.base64_encode(common.hex_to_bytes(expect_measurement)))

    -- spdm释放
    spdm_c:free()
end